from multiprocessing import Process
from fastapi import FastAPI, Request
from frame_transfer import yolo_frame, push_frame
from rtc_handler import rtc_frame
import queue

# 创建FastAPI应用实例
app = FastAPI()


# 定义一个函数，用于启动视频处理相关的进程
def start_video_processing(url, rtmp_url):
    # 创建一个队列用于存储从服务器接收的视频帧
    rtc_queue = queue.Queue(maxsize=10000)
    # 创建一个队列用于存储检测结果的视频帧
    yolo_queue = queue.Queue(maxsize=10000)

    import threading
    # 创建一个线程，用于从服务器接收视频帧
    producer_thread = threading.Thread(target=rtc_frame, args=(url, rtc_queue))
    # 创建一个线程，用于将原始视频帧队列中的帧转移到检测结果队列
    transfer_thread = threading.Thread(target=yolo_frame, args=(rtc_queue, yolo_queue))
    # 创建一个线程，用于从检测队列结果中获取帧并推送到RTMP服务器
    consumer_thread = threading.Thread(target=push_frame, args=(yolo_queue, rtmp_url))

    # 启动所有线程
    producer_thread.start()
    transfer_thread.start()
    consumer_thread.start()

    # 等待所有线程完成
    producer_thread.join()
    transfer_thread.join()
    consumer_thread.join()


@app.post("/start_video")
async def start_video(request: Request):
    # 从请求中获取JSON数据
    data = await request.json()
    url = data.get("url")
    rtmp_url = data.get("rtmp_url")
    # 启动一个新的进程来处理视频
    p = Process(target=start_video_processing, args=(url, rtmp_url))
    p.start()
    return {"message": "任务已启动"}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
